import tkinter as tk
from tkinter import filedialog, messagebox, ttk
from omegaconf import OmegaConf
from PIL import Image, ImageTk
from inference import inference_factory
import glob
import os


class Tooltip:
    """工具提示类"""

    def __init__(self, widget, text):
        self.widget = widget
        self.text = text
        self.tooltip_window = None
        self.widget.bind('<Enter>', self.show_tooltip)
        self.widget.bind('<Leave>', self.hide_tooltip)

    def show_tooltip(self, event=None):
        if self.tooltip_window is not None:
            return
        x = self.widget.winfo_rootx() + 20
        y = self.widget.winfo_rooty() + 20
        self.tooltip_window = tk.Toplevel(self.widget)
        self.tooltip_window.wm_overrideredirect(True)
        self.tooltip_window.wm_geometry(f"+{x}+{y}")
        label = tk.Label(self.tooltip_window, text=self.text, background='lightyellow', padx=5, pady=5)
        label.pack()

    def hide_tooltip(self, event=None):
        if self.tooltip_window:
            self.tooltip_window.destroy()
            self.tooltip_window = None


class OnnxModelApp:
    def __init__(self, root):
        self.root = root
        self.root.title("ONNX Model Inference GUI")
        self.config = OmegaConf.load("config.yaml")
        self.application = inference_factory(self.config)

        self.file_paths = []
        self.cur_idx = 0
        self.save_path = None
        self.tta = False

        # 按钮行
        self.button_frame = tk.Frame(self.root)
        self.button_frame.pack(pady=10)

        self.load_image_button = tk.Button(
            self.button_frame,
            text="加载图片",
            command=self.load_image,
            bg="#BAC8D3",
            fg="black",
            font=("宋体", 12, "bold"),
            height=1,
            width=10,
            relief=tk.RAISED
        )
        self.load_image_button.pack(side=tk.LEFT, padx=50)

        self.load_folder_button = tk.Button(self.button_frame, text="扫描文件夹", command=self.load_folder,
                                            bg="#BAC8D3",
                                            fg="black",
                                            font=("宋体", 12, "bold"),
                                            height=1,
                                            width=10,
                                            relief=tk.RAISED)
        self.load_folder_button.pack(side=tk.LEFT, padx=50)

        self.load_video_button = tk.Button(self.button_frame, text="加载视频流", command=self.load_video, bg="#BAC8D3",
                                           fg="black",
                                           font=("宋体", 12, "bold"),
                                           height=1,
                                           width=10,
                                           relief=tk.RAISED)
        self.load_video_button.pack(side=tk.LEFT, padx=50)

        self.start_inference_button = tk.Button(self.button_frame, text="开始推理", command=self.start_inference,
                                                bg="#D5E8D4",
                                                fg="black",
                                                font=("宋体", 12, "bold"),
                                                height=1,
                                                width=10,
                                                relief=tk.RAISED)
        self.start_inference_button.pack(side=tk.LEFT, padx=50)

        self.stop_inference_button = tk.Button(self.button_frame, text="终止推理", command=self.stop_inference,
                                               bg="#F8CECC",
                                               fg="black",
                                               font=("宋体", 12, "bold"),
                                               height=1,
                                               width=10,
                                               relief=tk.RAISED)
        self.stop_inference_button.pack(side=tk.LEFT, padx=50)

        self.save_button = tk.Button(self.button_frame, text="保存位置", command=self.get_save_path,
                                     bg="#E1D5E7",
                                     fg="black",
                                     font=("宋体", 12, "bold"),
                                     height=1,
                                     width=10,
                                     relief=tk.RAISED)
        self.save_button.pack(side=tk.RIGHT, padx=50)

        # 第二行
        self.settings_frame = tk.Frame(self.root)
        self.settings_frame.pack(pady=10)

        self.tta_var = tk.IntVar()
        cb = tk.Checkbutton(self.settings_frame, text='TTA', fg="black", font=("宋体", 12, "bold"),
                            variable=self.tta_var,
                            command=self.get_tta)
        cb.pack(side=tk.LEFT, padx=50)
        Tooltip(cb, '会对图像水平垂直翻转等操作，多种预测取均值，增加准确率(会增加推理时间)')

        t1 = tk.Frame(self.settings_frame)
        t1.pack(side=tk.LEFT, padx=50)

        tk.Label(t1, text='运行设备', fg="black", font=("宋体", 12, "bold")).pack(side=tk.LEFT)
        self.model_selection = ttk.Combobox(t1, values=["CPU", "GPU"], width=10)
        self.model_selection.set("CPU")
        self.model_selection.pack(side=tk.LEFT)

        t2 = tk.Frame(self.settings_frame)
        t2.pack(side=tk.LEFT, padx=50)
        self.interval_label = tk.Label(t2, text="推理时间间隔（毫秒）", fg="black", font=("宋体", 12, "bold"), )
        self.interval_label.pack(side=tk.LEFT)

        # 使用Spinbox创建可调节的推理时间间隔
        self.interval_spinbox = tk.Spinbox(t2, from_=0, to=1000, increment=1, width=10)
        self.interval_spinbox.pack(side=tk.LEFT)
        self.interval_spinbox.delete(0, tk.END)  # 清空默认文本
        self.interval_spinbox.insert(0, "50")  # 设置默认值为 1.0 秒

        # 图像展示区域
        self.image_frame = tk.Frame(self.root, width=1200, height=400)
        self.image_frame.pack_propagate(False)
        self.image_frame.pack(pady=10)

        self.image_label = tk.Label(self.image_frame, width=600, height=400, bg="#A1A1A1", relief=tk.RIDGE)
        self.image_label.pack(side=tk.LEFT)

        self.result_label = tk.Label(self.image_frame, width=600, height=400, bg="#A1A1A1", relief=tk.RIDGE)
        self.result_label.pack(side=tk.RIGHT)

        self.is_inference_running = False  # 控制推理状态

        self.log_frame = tk.Frame(self.root)
        self.log_frame.pack(pady=0)
        self.text_box = tk.Text(self.log_frame, width=167, height=7)  # 300x100 像素大约是 40x5 字符
        self.text_box.pack(side=tk.LEFT, fill=tk.BOTH, expand=True)

        # 创建滚动条
        self.scrollbar = tk.Scrollbar(self.log_frame, command=self.text_box.yview)
        self.scrollbar.pack(side=tk.RIGHT, fill=tk.Y)
        # 将滚动条与文本框关联
        self.text_box.config(yscrollcommand=self.scrollbar.set)

        self.info = tk.Frame(self.root)
        self.info.pack(side=tk.LEFT, pady=10)
        t3 = tk.Label(self.info, text='所展示图片会与真实大小有出入',
                      bg="#B46504",
                      fg="black",
                      font=("宋体", 12, "bold"),
                      height=1,
                      width=1200)
        t3.pack(side=tk.LEFT)

    def log_info(self, msg):
        self.text_box.insert(tk.END, f"{msg}\n")

    def load_image(self):
        file_path = filedialog.askopenfilename(filetypes=[("Image Files", "*.png;*.jpg;*.jpeg;*.bmp")])
        if file_path:
            self.file_paths.append(file_path)
            image = Image.open(file_path)
            image = image.resize((600, 400))  # 调整图像大小
            tk_image = ImageTk.PhotoImage(image)
            self.image_label.config(image=tk_image)
            self.image_label.image = tk_image  # 保持对图像的引用
            self.log_info(f'已加载图像: {file_path}')

    def load_folder(self):
        folder_path = filedialog.askdirectory()
        if folder_path:
            img_files = glob.glob(os.path.join(folder_path, "*.jpg")) + \
                        glob.glob(os.path.join(folder_path, "*.jpeg")) + \
                        glob.glob(os.path.join(folder_path, "*.bmp")) + \
                        glob.glob(os.path.join(folder_path, "*.png"))
            self.file_paths.extend(img_files)
            image = Image.open(self.file_paths[0])
            image = image.resize((600, 400))  # 调整图像大小
            tk_image = ImageTk.PhotoImage(image)
            self.image_label.config(image=tk_image)
            self.image_label.image = tk_image  # 保持对图像的引用
            self.log_info(f'已扫描文件夹: {folder_path}')
            self.log_info(f'已加载图像: {self.file_paths[0]}')

    def get_save_path(self):
        folder_path = filedialog.askdirectory()
        if folder_path:
            self.save_path = folder_path
            self.log_info(f'已选择保存路径: {folder_path}')

    def load_video(self):
        # 这里可以实现加载视频流的功能
        messagebox.showinfo("信息", "加载视频流功能尚未实现。")

    def start_inference(self):
        if self.save_path is None:
            messagebox.showerror('错误', '请选择保存位置')
            return
        if len(self.file_paths) == 0:
            messagebox.showerror('错误', '请上传样本')
            return
        if not self.is_inference_running:
            self.is_inference_running = True
            self.run_inference()

    def run_inference(self):
        if self.is_inference_running:
            self.log_info(f'开始推理样本`{self.file_paths[self.cur_idx]}`')
            if self.tta:
                self.application.load_img(self.file_paths[self.cur_idx]).tta_predict().save_res(self.save_path)
            else:
                self.application.load_img(self.file_paths[self.cur_idx]).sample_predict().save_res(self.save_path)
            if self.config.task_type == 'segment':
                output = Image.fromarray(self.application.res).resize((600, 400))  # 调整图像大小
                tk_image = ImageTk.PhotoImage(output)
                self.result_label.config(image=tk_image)
                self.result_label.image = tk_image  # 保持对图像的引用
            elif self.config.task_type == 'classify':
                pass
            elif self.config.task_type == 'unsupervised':
                pass
            self.log_info(f'已完成样本`{self.file_paths[self.cur_idx]}`的推理，结果已保存')

            if self.cur_idx < len(self.file_paths) - 1:
                self.cur_idx += 1
                image = Image.open(self.file_paths[self.cur_idx])
                image = image.resize((600, 400))  # 调整图像大小
                tk_image = ImageTk.PhotoImage(image)
                self.image_label.config(image=tk_image)
                self.image_label.image = tk_image  # 保持对图像的引用
                self.result_label.config(image=None)
                self.result_label.image = None
                self.log_info(f'已加载图像: {self.file_paths[self.cur_idx]}')
            else:
                messagebox.showinfo("成功", "推理完毕！")
                self.is_inference_running = False
                return
            try:
                interval = float(self.interval_spinbox.get())
            except ValueError:
                interval = 50  # 默认间隔为50毫秒

            self.root.after(int(interval / 1000), self.run_inference)  # 递归调用以实现间隔推理

    def stop_inference(self):
        self.is_inference_running = False

    def get_tta(self):
        self.tta = bool(self.tta_var.get())


if __name__ == "__main__":
    # 创建主窗口
    root = tk.Tk()
    app = OnnxModelApp(root)

    # 设置窗口大小
    window_width = 1200
    window_height = 650
    root.geometry(f"{window_width}x{window_height}")

    # 获取屏幕宽度和高度
    screen_width = root.winfo_screenwidth()
    screen_height = root.winfo_screenheight()

    # 计算窗口位置
    x = (screen_width // 2) - (window_width // 2)
    y = (screen_height // 2) - (window_height // 2)

    # 设置窗口显示位置
    root.geometry(f"{window_width}x{window_height}+{x}+{y}")

    # 运行主循环
    root.mainloop()
